#include "FFmpegWrapper.h"
#include <iostream>
#include <chrono>

FFmpegWrapper::FFmpegWrapper()
{
    avformat_network_init();
    frame_ = av_frame_alloc();
    packet_ = av_packet_alloc();
    sws_ctx_ = nullptr;
    max_reconnect_attempts_ = -1; // 增加最大重连次数
    current_reconnect_attempts_ = 0;
    reconnecting_ = false;
    reconnect_interval_ = 3; // 增加重连间隔时间

    // 保留原始连接参数-解码器
    av_dict_set(&opts, "rtsp_transport", "tcp", 0);
    av_dict_set(&opts, "stimeout", "5000000", 0);
    av_dict_set(&opts, "qmin", "10", 0);
    av_dict_set(&opts, "qmax", "51", 0);
}

FFmpegWrapper::~FFmpegWrapper()
{
    close();
    av_frame_free(&frame_);
    av_packet_free(&packet_);
    // 解码参数
    if (opts)
    {
        av_dict_free(&opts);
    }
}

bool FFmpegWrapper::open(const std::string &url)
{
    url_ = url;

    bool success = initDecoder(opts);

    if (!success)
    {
        return false;
    }

    running_ = true;
    decode_thread_ = std::thread(&FFmpegWrapper::decodeThreadFunc, this);

    // 启动业务处理线程
    process_running_ = true;
    process_thread_ = std::thread(&FFmpegWrapper::processThreadFunc, this);

    // 初始化编码推流
    if (!output_url_.empty())
    {
        if (!initEncoder())
        {
            std::cerr << "Failed to initialize encoder" << std::endl;
            return false;
        }
        encoding_ = true;
        streaming_ = true;
        encode_thread_ = std::thread(&FFmpegWrapper::encodeThreadFunc, this);
        stream_thread_ = std::thread(&FFmpegWrapper::streamThreadFunc, this);
    }

    return true;
}

void FFmpegWrapper::setStreamParams(const std::string &output_url, int bitrate, int fps)
{
    output_url_ = output_url;
    bitrate_ = bitrate;
    fps_ = fps;
}

void FFmpegWrapper::close()
{
    running_ = false;
    process_running_ = false;
    encoding_ = false;
    streaming_ = false;

    if (decode_thread_.joinable())
    {
        decode_thread_.join();
    }

    if (process_thread_.joinable())
    {
        process_thread_.join();
    }

    if (encode_thread_.joinable())
    {
        encode_thread_.join();
    }

    if (stream_thread_.joinable())
    {
        stream_thread_.join();
    }

    // 释放编码相关资源
    if (enc_ctx_)
    {
        avcodec_free_context(&enc_ctx_);
        enc_ctx_ = nullptr;
    }

    if (ofmt_ctx_)
    {
        if (!(ofmt_ctx_->oformat->flags & AVFMT_NOFILE))
        {
            avio_closep(&ofmt_ctx_->pb);
        }
        avformat_free_context(ofmt_ctx_);
        ofmt_ctx_ = nullptr;
    }

    cleanup();
}

bool FFmpegWrapper::initDecoder(AVDictionary *&opts)
{
    format_ctx_ = avformat_alloc_context();
    if (!format_ctx_)
    {
        std::cerr << "Failed to allocate format context" << std::endl;
        return false;
    }

    // 添加调试信息
    std::cout << "Opening input: " << url_ << std::endl;
    int ret = avformat_open_input(&format_ctx_, url_.c_str(), nullptr, &opts);
    if (ret != 0)
    {
        char errbuf[AV_ERROR_MAX_STRING_SIZE];
        av_strerror(ret, errbuf, sizeof(errbuf));
        std::cerr << "Failed to open input: " << url_ << ", error: " << errbuf << std::endl;
        return false;
    }

    if (avformat_find_stream_info(format_ctx_, nullptr) < 0)
    {
        std::cerr << "Failed to find stream info" << std::endl;
        return false;
    }

    video_stream_index_ = -1;
    for (unsigned int i = 0; i < format_ctx_->nb_streams; i++)
    {
        if (format_ctx_->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            video_stream_index_ = i;
            break;
        }
    }

    if (video_stream_index_ == -1)
    {
        std::cerr << "No video stream found" << std::endl;
        return false;
    }

    AVCodecParameters *codec_params = format_ctx_->streams[video_stream_index_]->codecpar;
    const AVCodec *codec = avcodec_find_decoder(codec_params->codec_id);
    if (!codec)
    {
        std::cerr << "Unsupported codec" << std::endl;
        return false;
    }

    codec_ctx_ = avcodec_alloc_context3(codec);
    if (!codec_ctx_)
    {
        std::cerr << "Failed to allocate codec context" << std::endl;
        return false;
    }

    if (avcodec_parameters_to_context(codec_ctx_, codec_params) < 0)
    {
        std::cerr << "Failed to copy codec parameters" << std::endl;
        return false;
    }

    // 设置解码器参数
    codec_ctx_->flags2 |= AV_CODEC_FLAG2_FAST; // 启用快速解码
    codec_ctx_->thread_count = 4;              // 使用4个线程解码

    AVDictionary *my_codec_open_opts = nullptr;
    if (avcodec_open2(codec_ctx_, codec, &my_codec_open_opts) < 0)
    {
        std::cerr << "Failed to open codec" << std::endl;
        return false;
    }

    width_ = codec_ctx_->width;
    height_ = codec_ctx_->height;

    // 初始化SWSContext用于图像格式转换
    sws_ctx_ = sws_getContext(width_, height_, codec_ctx_->pix_fmt,
                              width_, height_, AV_PIX_FMT_BGR24,
                              SWS_BILINEAR, nullptr, nullptr, nullptr);

    if (!sws_ctx_)
    {
        std::cerr << "Failed to initialize SWS context" << std::endl;
        return false;
    }

    if (my_codec_open_opts)
    {
        av_dict_free(&my_codec_open_opts);
    }

    return true;
}

void FFmpegWrapper::decodeThreadFunc()
{
    std::cout << "Decode thread started" << std::endl;
    while (running_)
    {
        if (!format_ctx_)
        {
            std::cerr << "Format context is null" << std::endl;
            break;
        }

        int ret = av_read_frame(format_ctx_, packet_);
        if (ret >= 0)
        {
            if (packet_->stream_index == video_stream_index_)
            {
                if (avcodec_send_packet(codec_ctx_, packet_) == 0)
                {
                    while (avcodec_receive_frame(codec_ctx_, frame_) == 0)
                    {
                        FrameData frame_data;
                        frame_data.frame = av_frame_alloc();
                        av_frame_ref(frame_data.frame, frame_);
                        frame_data.pts = frame_->pts;
                        frame_queue_.push(frame_data);
                    }
                }
            }
            av_packet_unref(packet_);
        }
        else
        {
            // 解码失败时停止编码和推流并释放资源
            encoding_ = false;
            streaming_ = false;
            if (encode_thread_.joinable())
            {
                encode_thread_.join();
            }

            if (stream_thread_.joinable())
            {
                stream_thread_.join();
            }

            // 释放编码相关资源
            if (enc_ctx_)
            {
                avcodec_free_context(&enc_ctx_);
                enc_ctx_ = nullptr;
            }

            if (ofmt_ctx_)
            {
                if (!(ofmt_ctx_->oformat->flags & AVFMT_NOFILE))
                {
                    avio_closep(&ofmt_ctx_->pb);
                }
                avformat_free_context(ofmt_ctx_);
                ofmt_ctx_ = nullptr;
            }

            // 清空编码队列
            encode_queue_.clear();

            while (running_)
            {
                if (reconnect() == -1)
                {
                    std::cerr << "Reconnect failed, stopping decode thread" << std::endl;
                    running_ = false;
                    break;
                }

                // 重连成功后重新启动编码推流
                if (!output_url_.empty())
                {
                    if (!initEncoder())
                    {
                        std::cerr << "Failed to reinitialize encoder" << std::endl;
                        continue;
                    }
                    encoding_ = true;
                    streaming_ = true;
                    
                    // 重新启动编码和推流线程
                    encode_thread_ = std::thread(&FFmpegWrapper::encodeThreadFunc, this);
                    stream_thread_ = std::thread(&FFmpegWrapper::streamThreadFunc, this);
                }
            }
        }
    }

    // 冲刷解码器
    avcodec_send_packet(codec_ctx_, nullptr);
    while (avcodec_receive_frame(codec_ctx_, frame_) == 0)
    {
        FrameData frame_data;
        frame_data.frame = av_frame_alloc();
        av_frame_ref(frame_data.frame, frame_);
        frame_data.pts = frame_->pts;
        frame_queue_.push(frame_data);
    }
}

void FFmpegWrapper::processThreadFunc()
{
    std::cout << "Process thread started" << std::endl;
    while (process_running_)
    {
        try
        {
            if (frame_queue_.empty())
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
                continue;
            }

            auto frame = frame_queue_.pop();
            if (!frame.frame)
            {
                std::cerr << "Received null frame" << std::endl;
                continue;
            }

            processFrame(frame);

            av_frame_unref(frame.frame);
            av_frame_free(&frame.frame);
        }
        catch (const std::exception &e)
        {
            std::cerr << "Process thread error: " << e.what() << std::endl;
        }
    }
}

void FFmpegWrapper::processFrame(const FrameData &frame)
{
    // 将解码后的帧放入编码队列
    FrameData encode_frame;
    encode_frame.frame = av_frame_alloc();
    av_frame_ref(encode_frame.frame, frame.frame);
    encode_frame.pts = frame.pts;
    encode_queue_.push(encode_frame);
}

void FFmpegWrapper::cleanup()
{
    if (sws_ctx_)
    {
        sws_freeContext(sws_ctx_);
        sws_ctx_ = nullptr;
    }

    if (codec_ctx_)
    {
        avcodec_free_context(&codec_ctx_);
        codec_ctx_ = nullptr;
    }

    if (format_ctx_)
    {
        avformat_close_input(&format_ctx_);
        format_ctx_ = nullptr;
    }
}

int FFmpegWrapper::reconnect()
{
    if (reconnecting_)
    {
        return 0;
    }
    reconnecting_ = true;

    if (max_reconnect_attempts_ > 0)
    {
        if (current_reconnect_attempts_ >= max_reconnect_attempts_)
        {
            std::cerr << "Max reconnect attempts reached: " << max_reconnect_attempts_ << " | exit reconnect" << std::endl;
            reconnecting_ = false;
            return -1;
        }

        if (reconnect_interval_ <= 0)
        {
            std::cerr << "reconnect interval: " << reconnect_interval_ << " | exit reconnect" << std::endl;
            reconnecting_ = false;
            return -1;
        }
    }

    current_reconnect_attempts_++;
    std::cout << "Reconnecting... Attempt " << current_reconnect_attempts_ << " of " << max_reconnect_attempts_ << std::endl;

    // 检查网络状态
    std::this_thread::sleep_for(std::chrono::seconds(reconnect_interval_));

    // 添加网络状态检查
    std::cout << "Checking network status..." << std::endl;
    int ping_result = system("ping -c 1 8.8.8.8 > /dev/null 2>&1");
    if (ping_result != 0)
    {
        std::cerr << "Network connection unavailable" << std::endl;
        return 0;
    }

    // 添加详细的资源释放日志
    std::cout << "Cleaning up resources before reconnecting..." << std::endl;
    cleanup();

    // 添加详细的初始化日志
    std::cout << "Initializing decoder for reconnection..." << std::endl;
    bool success = initDecoder(opts);

    if (success)
    {
        std::cout << "Reconnection successful" << std::endl;
        current_reconnect_attempts_ = 0;
        reconnecting_ = false;
        return -1;
    }

    std::cerr << "Reconnection failed" << std::endl;
    reconnecting_ = false;
    return 0;
}

bool FFmpegWrapper::initEncoder()
{
    // 初始化输出格式上下文
    avformat_alloc_output_context2(&ofmt_ctx_, nullptr, "flv", output_url_.c_str());
    if (!ofmt_ctx_)
    {
        std::cerr << "Could not create output context" << std::endl;
        return false;
    }

    // 查找H264编码器
    const AVCodec *codec = avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec)
    {
        std::cerr << "Codec not found" << std::endl;
        return false;
    }

    // 封装格式流
    video_st_ = avformat_new_stream(ofmt_ctx_, codec);
    if (!video_st_)
    {
        std::cerr << "Could not allocate video stream" << std::endl;
        return false;
    }

    // 创建编码器上下文
    enc_ctx_ = avcodec_alloc_context3(codec);
    if (!enc_ctx_)
    {
        std::cerr << "Could not allocate video codec context" << std::endl;
        return false;
    }

    // 设置编码参数
    enc_ctx_->codec_id = AV_CODEC_ID_H264;
    enc_ctx_->bit_rate = bitrate_;
    enc_ctx_->width = width_;
    enc_ctx_->height = height_;
    enc_ctx_->time_base = {1, fps_};
    enc_ctx_->framerate = {fps_, 1};
    enc_ctx_->gop_size = fps_; // 关键帧间隔
    enc_ctx_->max_b_frames = 1;
    enc_ctx_->pix_fmt = AV_PIX_FMT_YUV420P;

    // 打开编码器
    if (avcodec_open2(enc_ctx_, codec, nullptr) < 0)
    {
        std::cerr << "Could not open codec" << std::endl;
        return false;
    }

    // 设置流参数
    video_st_->time_base = enc_ctx_->time_base;
    if (avcodec_parameters_from_context(video_st_->codecpar, enc_ctx_) < 0)
    {
        std::cerr << "Could not copy codec parameters" << std::endl;
        return false;
    }

    // 打开输出文件
    if (!(ofmt_ctx_->oformat->flags & AVFMT_NOFILE))
    {
        if (avio_open(&ofmt_ctx_->pb, output_url_.c_str(), AVIO_FLAG_WRITE) < 0)
        {
            std::cerr << "Could not open output URL" << std::endl;
            return false;
        }
    }

    // 写文件头
    if (avformat_write_header(ofmt_ctx_, nullptr) < 0)
    {
        std::cerr << "Error occurred when writing header" << std::endl;
        return false;
    }

    return true;
}

void FFmpegWrapper::encodeThreadFunc()
{
    std::cout << "Encode thread started" << std::endl;
    while (encoding_)
    {
        try
        {
            if (encode_queue_.empty())
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
                continue;
            }

            auto frame = encode_queue_.pop();
            if (!frame.frame)
            {
                std::cerr << "Received null frame" << std::endl;
                continue;
            }

            // 转换帧格式为YUV420P
            AVFrame *yuv_frame = av_frame_alloc();
            yuv_frame->format = AV_PIX_FMT_YUV420P;
            yuv_frame->width = width_;
            yuv_frame->height = height_;
            av_frame_get_buffer(yuv_frame, 0);

            sws_scale(sws_ctx_, frame.frame->data, frame.frame->linesize, 0,
                      height_, yuv_frame->data, yuv_frame->linesize);

            // 编码帧
            yuv_frame->pts = frame.pts;
            if (avcodec_send_frame(enc_ctx_, yuv_frame) < 0)
            {
                std::cerr << "Error sending frame to encoder" << std::endl;
            }

            av_frame_free(&yuv_frame);
            av_frame_free(&frame.frame);
        }
        catch (const std::exception &e)
        {
            std::cerr << "Encode thread error: " << e.what() << std::endl;
        }
    }
}

void FFmpegWrapper::streamThreadFunc()
{
    std::cout << "Stream thread started" << std::endl;
    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = nullptr;
    pkt.size = 0;

    while (streaming_)
    {
        try
        {
            // 接收编码后的包
            int ret = avcodec_receive_packet(enc_ctx_, &pkt);
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            {
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
                continue;
            }
            else if (ret < 0)
            {
                std::cerr << "Error during encoding" << std::endl;
                continue;
            }

            // 写包
            pkt.stream_index = video_st_->index;
            av_packet_rescale_ts(&pkt, enc_ctx_->time_base, video_st_->time_base);
            if (av_interleaved_write_frame(ofmt_ctx_, &pkt) < 0)
            {
                std::cerr << "Error while writing video frame" << std::endl;
            }

            av_packet_unref(&pkt);
        }
        catch (const std::exception &e)
        {
            std::cerr << "Stream thread error: " << e.what() << std::endl;
        }
    }

    // 写文件尾
    av_write_trailer(ofmt_ctx_);
    av_packet_unref(&pkt);
}

bool FFmpegWrapper::isRunning() const
{
    return running_;
}

FrameData FFmpegWrapper::getFrame()
{
    return frame_queue_.pop();
}

int FFmpegWrapper::getWidth() const
{
    return width_;
}

int FFmpegWrapper::getHeight() const
{
    return height_;
}
